from configs.go2_constraint_him import Go2ConstraintHimRoughCfg, Go2ConstraintHimRoughCfgPPO
import cv2
import os
from isaacgym import gymtorch, gymapi, gymutil

from envs import LeggedRobot
from modules import *
from utils import  get_args, export_policy_as_jit, task_registry, Logger
from configs import *
from utils.helpers import class_to_dict
from utils.task_registry import task_registry
import numpy as np
import torch
from global_config import ROOT_DIR

from PIL import Image as im
from modules.actor_critic import ActorCriticRMA

def delete_files_in_directory(directory_path):
   try:
     files = os.listdir(directory_path)
     for file in files:
       file_path = os.path.join(directory_path, file)
       if os.path.isfile(file_path):
         os.remove(file_path)
     print("All files deleted successfully.")
   except OSError:
     print("Error occurred while deleting files.")

def play(args):
    # ---- 构造环境 ----
    env_cfg, train_cfg = task_registry.get_cfgs(name=args.task)
    env_cfg.env.num_envs = 1
    env_cfg.terrain.curriculum = False
    env_cfg.noise.add_noise = False
    env_cfg.domain_rand.randomize_friction = False
    env_cfg.domain_rand.randomize_base_com = False
    env_cfg.domain_rand.randomize_base_mass = False
    env_cfg.domain_rand.randomize_motor = False
    env_cfg.domain_rand.randomize_lag_timesteps = False
    env_cfg.domain_rand.randomize_restitution = False
    env_cfg.domain_rand.disturbance = False
    env_cfg.domain_rand.randomize_kpkd = False
    env_cfg.control.use_filter = True

    env, _ = task_registry.make_env(name=args.task, args=args, env_cfg=env_cfg)

    _root_states = gymtorch.wrap_tensor(env.gym.acquire_actor_root_state_tensor(env.sim))   
    env.gym.refresh_actor_root_state_tensor(env.sim)
    _root_states[:, 2] += 0.25       
    _root_states[:, 7:13] = 0        
    env.gym.set_actor_root_state_tensor(env.sim,gymtorch.unwrap_tensor(_root_states))

    obs = env.get_observations()

    # ---- 加载 checkpoint ----
    ckpt_path = '/home/rog/LocomotionWithNP3O/logs/rough_go2_constraint/May26_21-33-19_test_barlowtwins/model_10000.pt'
    ckpt = torch.load(ckpt_path, map_location=env.device, weights_only=True)

    # ---- 构造所有模块 ----
    policy_cfg_dict = class_to_dict(train_cfg.policy)
    runner_cfg_dict = class_to_dict(train_cfg.runner)
    actor_critic_class = eval(runner_cfg_dict["policy_class_name"])
    actor_critic = actor_critic_class(
        env.cfg.env.n_proprio,
        env.cfg.env.n_scan,
        env.num_obs,
        env.cfg.env.n_priv_latent,
        env.cfg.env.history_len,
        env.num_actions,
        **policy_cfg_dict
    ).to(env.device).half()
    actor_critic.load_state_dict(ckpt['model_state_dict'])
    actor_critic.eval()

    # ---- 相机设置 ----
    camera_local_transform = gymapi.Transform()
    camera_local_transform.p = gymapi.Vec3(-0.5, -1, 0.1)
    camera_local_transform.r = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 0, 1), np.deg2rad(90))
    camera_props = gymapi.CameraProperties()
    camera_props.width = 512
    camera_props.height = 512
    cam_handle = env.gym.create_camera_sensor(env.envs[0], camera_props)
    body_handle = env.gym.get_actor_rigid_body_handle(env.envs[0], env.actor_handles[0], 0)
    env.gym.attach_camera_to_body(cam_handle, env.envs[0], body_handle, camera_local_transform, gymapi.FOLLOW_TRANSFORM)

    action_rate = 0
    z_vel = 0
    xy_vel = 0
    feet_air_time = 0

    # ---- 主循环 ----
    video = None
    num_frames = int(20 / env.dt)
    for i in range(num_frames):
        action_rate += torch.sum(torch.abs(env.last_actions - env.actions),dim=1)
        z_vel += torch.square(env.base_lin_vel[:, 2])
        xy_vel += torch.sum(torch.square(env.base_ang_vel[:, :2]), dim=1)

        env.commands[:,0] = 1
        env.commands[:,1] = 0
        env.commands[:,2] = 0
        env.commands[:,3] = 0
        actions = actor_critic.act_teacher(obs.half())
        # actions = torch.clamp(actions,-1.2,1.2)

        obs, privileged_obs, rewards,costs,dones, infos = env.step(actions)
        env.gym.step_graphics(env.sim) # required to render in headless mode
        env.gym.render_all_camera_sensors(env.sim)

        # ---- 渲染 ----
        if not env.headless:
            env.gym.draw_viewer(env.viewer, env.sim, True)
        env.gym.step_graphics(env.sim)
        env.gym.render_all_camera_sensors(env.sim)

        # ---- 录屏 ----
        if RECORD_FRAMES:
            img = env.gym.get_camera_image(env.sim, env.envs[0], cam_handle, gymapi.IMAGE_COLOR).reshape(512, 512, 4)[:, :, :3]
            if video is None:
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                jit_dir = os.path.join(ROOT_DIR, 'logs', 'play')
                os.makedirs(jit_dir, exist_ok=True)
                video = cv2.VideoWriter(os.path.join(jit_dir, 'play.mp4'), fourcc, int(1 / env.dt), (512, 512))
            video.write(img[..., ::-1])

    if video is not None:
        video.release()

if __name__ == '__main__':
    task_registry.register("go2N3poHim", LeggedRobot,Go2ConstraintHimRoughCfg(),Go2ConstraintHimRoughCfgPPO())
   
    RECORD_FRAMES = True
    args = get_args()
    
    play(args)
